#include <fstream>
#include <iostream>
#include <cstdlib>
#include <string>
#include <cmath>
#include <mpi.h>

#include "CmdLineArgs.h"
#include "MPIIO.h"
#include "PWaveProblem.h"
#include "PWaveData3D.h"
#include "PWaveStencilMath.h"

#define INDEX(x, y, z)  ((x) + (__int64) ldx * ((y) + (__int64) ldy * (z)))

#define DOUBLE_EPSILON 1e-15
#define M_PI 3.14159265358979323846


void PWaveProblem::init(const CmdLineArgs* cmd) {
    std::ifstream fin(cmd->in_file.c_str());

    if(!fin.good()) {
        std::cout << "No such file: " << cmd->in_file << " !\n";
        exit(1);
    }

    fin >> NGP_X >> NGP_Y >> NGP_Z;
    fin >> PML_XN >> PML_XP >> PML_YN >> PML_YP >> PML_ZN >> PML_ZP;
    fin >> halo;

    fin >> this->NUM_BLK_Y >> this->NUM_BLK_Z;

    fin >> end_time_steps >> dt;
    fin >> realX >> realY >> realZ;

    
    fin >> model_src;
    if(model_src == "built-in") {
        use_outer_model = false;
        fin >> c1 >> c2 >> h1 >> h2;
        h1 += PML_XN;
        h2 += PML_XP;
    } else if(model_src == "outer") {
        use_outer_model = true;
        model_src = cmd->c_model;
    } else {
        std::cout << "unrecognized velocity model source: " << model_src << "\n";
        exit(1);
    }

    fin >> num_src;
    src_info.reserve(num_src);

    for(int i = 0; i < num_src; i++) {
        int temp_x, temp_y, temp_z;
        double temp_f, temp_amp, temp_t0;
        fin >> temp_x >> temp_y >> temp_z >> temp_f >> temp_amp >> temp_t0;

        struct GlobalSrcInfo temp;
        temp.x = temp_x + PML_XN;
        temp.y = temp_y + PML_YN;
        temp.z = temp_z + PML_ZN;
        temp.f = temp_f;
        temp.amp = temp_amp;
        temp.t0 = temp_t0;

        src_info.push_back(temp);
    }

    fin >> B;

    for(int i = 0; i < 3; i++) {
        std::string temp;
        fin >> temp;
        if(temp == "x" || temp == "X") {
            fin >> frames_freq_x >> target_x;
            target_x += PML_XN;
            if(frames_freq_x <= 0)
                output_iter_frames_x = false;
            else
                output_iter_frames_x = true;
        } else if(temp == "y" || temp == "Y") {
            fin >> frames_freq_y >> target_y;
            target_y += PML_YN;
            if(frames_freq_y <= 0)
                output_iter_frames_y = false;
            else
                output_iter_frames_y = true;
        } else if(temp == "z" || temp == "Z") {
            fin >> frames_freq_z >> target_z;
            target_z += PML_ZN;
            if(frames_freq_z <= 0)
                output_iter_frames_z = false;
            else
                output_iter_frames_z = true;
        } else {
            std::cout << "unrecognized input parameter: " << temp << "\n";
            exit(1);
        }
    }

    fin >> num_output_pts;
    output_pts_info.reserve(num_output_pts);
    for(int i = 0; i < num_output_pts; i++) {
        int temp_x, temp_y, temp_z;
        fin >> temp_x >> temp_y >> temp_z;
        struct GlobalSrcInfo temp;
        temp.x = temp_x + PML_XN;
        temp.y = temp_y + PML_YN;
        temp.z = temp_z + PML_ZN;
        temp.f = 0;
        temp.amp = 0;
        temp.t0 = 0;
        output_pts_info.push_back(temp);
    }


    fin >> identifier;

    raw = cmd->raw;
    out_final3D = cmd->out_final_3D;

    out_prefix = cmd->out_dir + identifier;

    //! fill in default params

    h = realX / (NGP_X - 1);   // temporarily use only X dim to calculate h, and we don't check whether it is consistent with Y and Z

    if(B < 0)       // use default value for B
        B = 1e-4;

    if(!use_outer_model) {
        double c_m = (c1 > c2) ? c1 : c2;
        if(dt < 0)      //use default configuration for dt, which is: c*dt/h=0.2
            dt = 0.2 * h / c_m;

        //! check input data
        if(c_m * dt / h > 0.6124)
            std::cout << "Warning: iteration might not converge, for cm * dt / h = " << c_m * dt / h << " > 0.6124.\n";
    }
}



void PWaveProblem::partition(Data3D* data3D) {
    int comm_siz, proc_id;
    MPI_Comm_size(MPI_COMM_WORLD, &comm_siz);
    MPI_Comm_rank(MPI_COMM_WORLD, &proc_id);

    data3D->halo = halo;

    data3D->x_size = PML_XN + NGP_X + PML_XP;
    data3D->y_size = PML_YN + NGP_Y + PML_YP;
    data3D->z_size = PML_ZN + NGP_Z + PML_ZP;


    //! then local (block-wise) grid info

    // divide the whole grid in y and z direction (a 2D-cut)
    int ldx, ldy, ldz;
    ldx = 2 * halo + data3D->x_size;
    data3D->ldx = ldx;
    int NUM_BLK_Y = this->NUM_BLK_Y;
    int NUM_BLK_Z = this->NUM_BLK_Z;

    int BLK_LOC_Y = proc_id % NUM_BLK_Y;
    int BLK_LOC_Z = proc_id / NUM_BLK_Y;

    data3D->BLK_SIZE_Y = data3D->y_size / NUM_BLK_Y;        // store the normal block size for each dimension
    data3D->BLK_SIZE_Z = data3D->z_size / NUM_BLK_Z;
    std::cout << "proc " << proc_id << ", BLK_LOC_Y = " << BLK_LOC_Y << ", BLK_LOC_Z = " << BLK_LOC_Z << "\n";


    // if [y|z]_size isn't a multiple of NUM_BLK_[Y|Z], the size of last block will be different from others
    if(BLK_LOC_Y == NUM_BLK_Y - 1)
        data3D->local_blk_size_y = data3D->y_size - BLK_LOC_Y * data3D->BLK_SIZE_Y;
    else
        data3D->local_blk_size_y = data3D->BLK_SIZE_Y;

    if(BLK_LOC_Z == NUM_BLK_Z - 1)
        data3D->local_blk_size_z = data3D->z_size - BLK_LOC_Z * data3D->BLK_SIZE_Z;
    else
        data3D->local_blk_size_z = data3D->BLK_SIZE_Z;

    ldy = data3D->local_blk_size_y + 2 * halo;
    ldz = data3D->local_blk_size_z + 2 * halo;

    data3D->ldy = ldy;
    data3D->ldz = ldz;

    data3D->blk_domain_size = (__int64)ldx * (__int64)ldy * (__int64)ldz;



    // record the position of wave src
    for(int i = 0; i < num_src; i++) {
        int abs_src_x = src_info[i].x;            // temporarily ignoring the existence of halo area
        int abs_src_y = src_info[i].y;
        int abs_src_z = src_info[i].z;
        int src_blk_y_id = abs_src_y / data3D->BLK_SIZE_Y;
        int src_blk_z_id = abs_src_z / data3D->BLK_SIZE_Z;

        // if this proc is responsible of the block where the src is located in, compute the pos of the src in the block
        if(BLK_LOC_Y == src_blk_y_id && BLK_LOC_Z == src_blk_z_id) {
            struct LocalSrcInfo temp;
            temp.inblk_src_loc = INDEX(abs_src_x + halo,              // assume no blocking in x dir
                                      (abs_src_y - src_blk_y_id * data3D->BLK_SIZE_Y) + halo,
                                      (abs_src_z - src_blk_z_id * data3D->BLK_SIZE_Z) + halo);
            temp.f = src_info[i].f;
            temp.amp = src_info[i].amp;
            temp.t0 = src_info[i].t0;

            inblk_src_info.push_back(temp);

            std::cout << "proc " << proc_id << " contains wave src " << i << " .\n";
        }
    }


    // record the position of output points:
    for(int i = 0; i < num_output_pts; i++) {
        int abs_out_x = output_pts_info[i].x;            // temporarily ignoring the existence of halo area
        int abs_out_y = output_pts_info[i].y;
        int abs_out_z = output_pts_info[i].z;
        int out_blk_y_id = abs_out_y / data3D->BLK_SIZE_Y;
        int out_blk_z_id = abs_out_z / data3D->BLK_SIZE_Z;

        // if this proc is responsible of the block where the output point is located in, compute the pos of the point in the block
        if(BLK_LOC_Y == out_blk_y_id && BLK_LOC_Z == out_blk_z_id) {
            char coord_str[40];
            sprintf(coord_str, "_out_pt_id%d_x%d_y%d_z%d.dat", i, abs_out_x - PML_XN, abs_out_y - PML_YN, abs_out_z - PML_ZN);
            std::string s = out_prefix + coord_str;
            FILE* fp = fopen(s.c_str(), "rb");
            if(fp) {
                fclose(fp);
                remove(s.c_str());
            }

            OutputPtInfo temp;
            temp.name = s;
            temp.pos = INDEX(abs_out_x + halo,              // assume no blocking in x dir
                            (abs_out_y - out_blk_y_id * data3D->BLK_SIZE_Y) + halo,
                            (abs_out_z - out_blk_z_id * data3D->BLK_SIZE_Z) + halo);

            inblk_output_pt_info.push_back(temp);

            std::cout << "proc " << proc_id << " contains output point " << i << " .\n";
        }
    }
}


void PWaveProblem::create_MPI_topo(MPI_Topo* topo) {
    MPI_TopoCart2D* topo2D = dynamic_cast<MPI_TopoCart2D*>(topo);
    topo2D->init();
    topo2D->NUM_BLK_Y = this->NUM_BLK_Y;
    topo2D->NUM_BLK_Z = this->NUM_BLK_Z;

    topo2D->get_self_loc();
    topo2D->get_neighbors();
}


void PWaveProblem::init_model(Data3D* data3D, const MPI_Topo* topo) {
    PWaveData3D* pwave_data = dynamic_cast<PWaveData3D*>(data3D);
    const MPI_TopoCart2D* topo2D = dynamic_cast<const MPI_TopoCart2D*>(topo);


    // set computing domain to zeros, and init non_pml_x_[min|max] = halo for default situation
    // this function should be called before q initialization,
    // because non_pml_x_[min|max] will be modified to their corresponding correct values in some blocks during q initialization
    pwave_data->init();

    PWaveStencilMath::compute_wave_src(pwave_data->compute_domain, inblk_src_info, 0);
    

    //! initialize c model into coeffc2_distr
    if(use_outer_model) {
        MPIIO input_agent;
        input_agent.read_yz_block(pwave_data->c2dtq_distr, model_src.c_str(), data3D, topo2D);
    } else {
        // initialize two-layer c model from .in file
        printf("\ninitialize two-layer c model from .in file.\n");
        
        for(int k = 0; k < data3D->local_blk_size_z; k++) {                     // already taken the halo area into account
            __int64 disp_k = (__int64) data3D->local_blk_size_y * k;

            for(int j = 0; j < data3D->local_blk_size_y; j++) {
                __int64 disp_jk = data3D->x_size * (j + disp_k);

                int i;
                for(i = 0; i < h1; i++)
                    pwave_data->c2dtq_distr[i + disp_jk] = c1;
                for(; i < h1 + h2; i++)
                    pwave_data->c2dtq_distr[i + disp_jk] = c2;
            }
        }
    }

    int ldx = data3D->ldx, ldy = data3D->ldy;
    __int64 blk_domain_size = data3D->blk_domain_size;

    // use coeffc2_distr to temporarily store c^2, since c^2 will be used in the initialization process of c^2dtq_distr
    for(int k = 0; k < data3D->local_blk_size_z; k++) {
        __int64 disp_k = (__int64) k * data3D->local_blk_size_y;

        for(int j = 0; j < data3D->local_blk_size_y; j++) {
            __int64 disp_jk = (disp_k + j) * data3D->x_size;

            for(int i = 0; i < data3D->x_size; i++) {
                double c_temp = pwave_data->c2dtq_distr[disp_jk + i];
                pwave_data->coeffc2_distr[INDEX(i + halo, j + halo, k + halo)] = c_temp * c_temp;
            }
        }
    }




    //! set absorbing layer value field Q * c^2 * dt
    for(__int64 i = 0; i < blk_domain_size; i++)
        pwave_data->c2dtq_distr[i] = 0;

    int y_delta = topo2D->BLK_LOC_Y * data3D->BLK_SIZE_Y;
    int z_delta = topo2D->BLK_LOC_Z * data3D->BLK_SIZE_Z;


    int x_1l = 0;
    int x_1r = PML_XN;
    int x_2l = NGP_X + PML_XN;
    int x_2r = NGP_X + PML_XN + PML_XP;

    int y_1l = -y_delta;
    int y_1r = PML_YN - y_delta;
    int y_2l = NGP_Y + PML_YN - y_delta;
    int y_2r = NGP_Y + PML_YN + PML_YP - y_delta;

    int z_1l = -z_delta;
    int z_1r = PML_ZN - z_delta;
    int z_2l = NGP_Z + PML_ZN - z_delta;
    int z_2r = NGP_Z + PML_ZN + PML_ZP - z_delta;

    double Bdt = B * dt;

    int non_pml_x_min = x_2r, non_pml_x_max = x_1l;
    int non_pml_y_min = y_2r, non_pml_y_max = y_1l;
    int non_pml_z_min = z_2r, non_pml_z_max = z_1l;
    for(int z_local = 0; z_local < data3D->local_blk_size_z; z_local++) {
        double z_term = 0;
        if((z_1l <= z_local) && (z_local < z_1r)) {
            double dz = z_1r - z_local;
            z_term = dz * dz / PML_ZN / PML_ZN;

        } else if((z_2l <= z_local) && (z_local < z_2r)) {
            double dz = z_local - z_2l + 1;
            z_term = dz * dz / PML_ZP / PML_ZP;
        }


        for(int y_local = 0; y_local < data3D->local_blk_size_y; y_local++) {
            double y_term = 0;
            if((y_1l <= y_local) && (y_local < y_1r)) {
                double dy = y_1r - y_local;
                y_term = dy * dy / PML_YN / PML_YN;

            } else if((y_2l <= y_local) && (y_local < y_2r)) {
                double dy = y_local - y_2l + 1;
                y_term = dy * dy / PML_YP / PML_YP;
            }

            for(int x_local = 0; x_local < x_2r; x_local++) {       // x direction is not divided
                double x_term = 0;
                if((x_1l <= x_local) && (x_local < x_1r)) {
                    double dx = x_1r - x_local;
                    x_term = dx * dx / PML_XN / PML_XN;
                } else if((x_2l <= x_local) && (x_local < x_2r)) {
                    double dx = x_local - x_2l + 1;
                    x_term = dx * dx / PML_XP / PML_XP;
                }


                double d = sqrt(x_term + y_term + z_term);
                if(d > 1.0)
                    d = 1.0;
                else if(-DOUBLE_EPSILON < d && d < DOUBLE_EPSILON) {    // in Q=0 zone
                    if(z_local < non_pml_z_min)
                        non_pml_z_min = z_local;
                    else if(z_local > non_pml_z_max)
                        non_pml_z_max = z_local;

                    if(y_local < non_pml_y_min)
                        non_pml_y_min = y_local;
                    else if(y_local > non_pml_y_max)
                        non_pml_y_max = y_local;

                    if(x_local < non_pml_x_min)
                        non_pml_x_min = x_local;
                    else if(x_local > non_pml_x_max)
                        non_pml_x_max = x_local;
                }

                pwave_data->c2dtq_distr[INDEX(x_local + halo, y_local + halo, z_local + halo)] =
                (1. - cos(M_PI / 2 * d)) * Bdt * pwave_data->coeffc2_distr[INDEX(x_local + halo, y_local + halo, z_local + halo)];
            }
        }
    }

    if(non_pml_y_min == 0 && non_pml_y_max == data3D->local_blk_size_y - 1 && non_pml_z_min == 0 && non_pml_z_max == data3D->local_blk_size_z - 1) {
        pwave_data->non_pml_x_min = halo + non_pml_x_min;
        pwave_data->non_pml_x_max = halo + non_pml_x_max + 1;
    }



    //! after initialization of c2dtq, proceed to the calculation of coeffc2
    double coeff = dt * dt / h / h;

    for(int k = 0; k < data3D->local_blk_size_z; k++)
        for(int j = 0; j < data3D->local_blk_size_y; j++)
            for(int i = 0; i < data3D->x_size; i++)
                pwave_data->coeffc2_distr[INDEX(i + halo, j + halo, k + halo)] *= coeff;
}


double* PWaveProblem::solve(Data3D* data3D, const MPI_Topo* topo) {
    int ldx = data3D->ldx, ldy = data3D->ldy, ldz = data3D->ldz;
    int local_blk_size_z = data3D->local_blk_size_z;
    __int64 blk_domain_size = data3D->blk_domain_size;

    PWaveData3D* pwave_data = dynamic_cast<PWaveData3D*>(data3D);
    const MPI_TopoCart2D* topo2D = dynamic_cast<const MPI_TopoCart2D*>(topo);

    double* area_1 = pwave_data->compute_domain;
    double* area_2 = pwave_data->compute_domain + blk_domain_size;
#ifdef __DEBUG__
    printf("before req\n");
#endif
    //! MPI statuses and requests declaration
    MPI_Request requests_odd[8];
    MPI_Request requests_even[8];
    MPI_Request* requests[2] = {requests_even, requests_odd};

    //! MPI vector type preparation
    MPI_Datatype XZ_Exchange;
    MPI_Type_vector(local_blk_size_z, 2 * ldx, (__int64)ldx * ldy, MPI_DOUBLE, &XZ_Exchange);
    MPI_Type_commit(&XZ_Exchange);


    //! prepare contents for non-blocking repeated communication
    MPI_Recv_init(area_1 + (__int64)ldx * ldy * halo, 1, XZ_Exchange, topo2D->NBR_YN, 0, MPI_COMM_WORLD, &requests_odd[0]);   // -y dir
    MPI_Recv_init(area_1, (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZN, 0, MPI_COMM_WORLD, &requests_odd[1]);            // -z dir
    MPI_Recv_init(area_1 + (__int64)ldx * ldy * (ldz - halo), (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZP, 0, MPI_COMM_WORLD, &requests_odd[2]);  // +z dir
    MPI_Recv_init(area_1 + (__int64)ldx * ldy * halo + ldx * (ldy - halo), 1, XZ_Exchange, topo2D->NBR_YP, 0, MPI_COMM_WORLD, &requests_odd[3]);   // +y dir

    MPI_Send_init(area_1 + (__int64)ldx * ldy * halo + ldx * halo, 1, XZ_Exchange, topo2D->NBR_YN, 0, MPI_COMM_WORLD, &requests_odd[4]);
    MPI_Send_init(area_1 + (__int64)ldx * ldy * halo, (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZN, 0, MPI_COMM_WORLD, &requests_odd[5]);
    MPI_Send_init(area_1 + (__int64)ldx * ldy * (ldz - 2 * halo), (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZP, 0, MPI_COMM_WORLD, &requests_odd[6]);
    MPI_Send_init(area_1 + (__int64)ldx * ldy * halo + (__int64)ldx * (ldy - 2 * halo), 1, XZ_Exchange, topo2D->NBR_YP, 0, MPI_COMM_WORLD, &requests_odd[7]);


    MPI_Recv_init(area_2 + (__int64)ldx * ldy * halo, 1, XZ_Exchange, topo2D->NBR_YN, 0, MPI_COMM_WORLD, &requests_even[0]);
    MPI_Recv_init(area_2, (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZN, 0, MPI_COMM_WORLD, &requests_even[1]);
    MPI_Recv_init(area_2 + (__int64)ldx * ldy * (ldz - halo), (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZP, 0, MPI_COMM_WORLD, &requests_even[2]);
    MPI_Recv_init(area_2 + (__int64)ldx * ldy * halo + ldx * (ldy - halo), 1, XZ_Exchange, topo2D->NBR_YP, 0, MPI_COMM_WORLD, &requests_even[3]);

    MPI_Send_init(area_2 + (__int64)ldx * ldy * halo + ldx * halo, 1, XZ_Exchange, topo2D->NBR_YN, 0, MPI_COMM_WORLD, &requests_even[4]);
    MPI_Send_init(area_2 + (__int64)ldx * ldy * halo, (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZN, 0, MPI_COMM_WORLD, &requests_even[5]);
    MPI_Send_init(area_2 + (__int64)ldx * ldy * (ldz - 2 * halo), (__int64)halo * ldx * ldy, MPI_DOUBLE, topo2D->NBR_ZP, 0, MPI_COMM_WORLD, &requests_even[6]);
    MPI_Send_init(area_2 + (__int64)ldx * ldy * halo + (__int64)ldx * (ldy - 2 * halo), 1, XZ_Exchange, topo2D->NBR_YP, 0, MPI_COMM_WORLD, &requests_even[7]);


    /* begin computing process */
    double* domains[2] = {area_1, area_2};  // to store values in t and t-dt time steps respectively

    int non_pml_min_x = pwave_data->non_pml_x_min;
    int non_pml_max_x = pwave_data->non_pml_x_max;


    MPIIO output_agent;
    double* output_buff = new double[inblk_output_pt_info.size() * output_agent.buff_size];
    // temporarily stores the data to be outputted, output per buff_size time steps


    double* coeffc2_dist = pwave_data->coeffc2_distr, *c2dtq_distr = pwave_data->c2dtq_distr;

    for(int time_step = 1; time_step <= end_time_steps; time_step++) {
#ifndef SIMPLE_ITER_OUT
        if(time_step % 100 == 0)
            printf("process started iteration %d!\n", time_step);
#endif
        double* src = domains[(time_step - 1) % 2];
        double* dst = domains[time_step % 2];


        // even if wave src appears in somebody's halo area, the initial value will be exchanged by this call
        MPI_Request* request_t = requests[time_step % 2];
        MPI_Startall(8, request_t);

        PWaveStencilMath::stencil_order_4_pml_core(dst, src, coeffc2_dist, c2dtq_distr, halo, ldx, ldy, ldz, non_pml_min_x, non_pml_max_x);

        int index_arrived;
        for(int j = 0; j < 8; j++) {
            MPI_Waitany(8, request_t, &index_arrived, MPI_STATUS_IGNORE);

            switch(index_arrived) {
                case 0:     // -y dir
                    PWaveStencilMath::stencil_order_4_pml_kern(dst, src, coeffc2_dist, c2dtq_distr,
                                                        halo, halo, 2 * halo,
                                                        ldx - halo, 2 * halo, ldz - 2 * halo, ldx, ldy);
                    break;

                case 1:     // -z dir
                    PWaveStencilMath::stencil_order_4_pml_kern(dst, src, coeffc2_dist, c2dtq_distr,
                                                        halo, 2 * halo, halo,
                                                        ldx - halo, ldy - 2 * halo, 2 * halo, ldx, ldy);
                    break;

                case 2:     // +z dir
                    PWaveStencilMath::stencil_order_4_pml_kern(dst, src,coeffc2_dist, c2dtq_distr,
                                                        halo, 2 * halo, ldz - 2 * halo,
                                                        ldx - halo, ldy - 2 * halo, ldz - halo, ldx, ldy);
                    break;

                case 3:     // +y dir
                    PWaveStencilMath::stencil_order_4_pml_kern(dst, src, coeffc2_dist, c2dtq_distr,
                                                        halo, ldy - 2 * halo, 2 * halo,
                                                        ldx - halo, ldy - halo, ldz - 2 * halo, ldx, ldy);
                    break;
            }
        }


        //! then process the four corners left
        PWaveStencilMath::stencil_order_4_pml_corner(dst, src, coeffc2_dist, c2dtq_distr, halo, ldx, ldy, ldz);



        // finally update src term
        PWaveStencilMath::compute_wave_src(dst, inblk_src_info, time_step * dt);

        // printf("finish compute source in iteration %d.\n", time_step);


        if(output_iter_frames_x && time_step % frames_freq_x == 0) {
            // printf("step=%d\tsave X slice data.\n", time_step);
            char time_space[30];
            sprintf(time_space, "_x%d_t%d.dat", target_x - PML_XN, time_step);

            output_agent.blk3Dslice_x(dst, out_prefix + time_space, raw, data3D, topo2D, target_x);
        }

        if(output_iter_frames_y && time_step % frames_freq_y == 0) {
            // printf("step=%d\tsave Y slice data.\n", time_step);
            char time_space[30];
            sprintf(time_space, "_y%d_t%d.dat", target_y - PML_YN, time_step);

            output_agent.blk3Dslice_y(dst, out_prefix + time_space, raw, data3D, topo2D, target_y);
        }

        if(output_iter_frames_z && time_step % frames_freq_z == 0) {
            // printf("step=%d\tsave Z slice data.\n", time_step);
            char time_space[30];
            sprintf(time_space, "_z%d_t%d.dat", target_z - PML_ZN, time_step);

            output_agent.blk3Dslice_z(dst, out_prefix + time_space, raw, data3D, topo2D, target_z);
        }

        output_agent.out_coord_val(dst, output_buff, inblk_output_pt_info);
    }
    output_agent.flush_buffer(output_buff, inblk_output_pt_info);

    if(out_final3D)
        output_agent.blk3D2arr(domains[end_time_steps % 2], out_prefix, raw, data3D, topo2D);

    delete[] output_buff;
    MPI_Request_free(requests_odd);
    MPI_Request_free(requests_even);

    return domains[end_time_steps % 2];
}


void PWaveProblem::free() {

}


void PWaveProblem::write_log(const Data3D* data3D, double init_abs_sec, double compute_abs_sec,
                             double output_abs_sec, double start_abs_sec) const {
    // write log for PWave problem; 2021.04.08, sun weitao
    std::ofstream logfile;
    int comm_siz;
    MPI_Comm_size(MPI_COMM_WORLD, &comm_siz);
    char file_log[255] = {""};
    

    sprintf(file_log, "%s_%dx%dx%d_np%d_timestep%d.log", out_prefix.c_str(), NGP_X, NGP_Y, NGP_Z, comm_siz, end_time_steps);
    logfile.open(file_log);

    logfile <<  "xsize =" << realX  << "\tysize =" << realY << "\t zsize =" << realZ << "\n";   // Goes to file
    logfile <<  "nx =" << NGP_X  << "\tny =" << NGP_Y << "\t nz =" << NGP_Z << "\n";   // Goes to file
    logfile <<  "ldx =" << data3D->ldx  << "\tldy =" << data3D->ldy << "\t ldz =" << data3D->ldz << "\n";   // Goes to file
    logfile <<  "np =" << comm_siz  << "\ttime step =" << end_time_steps << "\n";   // Goes to file
    logfile <<  "time_init =\t" << init_abs_sec - start_abs_sec << " s\ntime_compute =\t" << compute_abs_sec - init_abs_sec
            << " s\ntime_output =\t" << output_abs_sec - compute_abs_sec << " s\n";   // Goes to file
    logfile.close();
}
